#include <iostream>
#include <string>
#include <cstring>
#include <cstdlib>
#include <unistd.h>
#include <arpa/inet.h>
#include <getopt.h>

#define BUFFER_SIZE 2048

// 解析IP和端口
bool parse_ip_port(const std::string &str, std::string &ip, int &port) {
    size_t colon_pos = str.find(':');
    if (colon_pos == std::string::npos) {
        std::cerr << "Invalid format: " << str << std::endl;
        return false;
    }
    ip = str.substr(0, colon_pos);
    port = std::atoi(str.substr(colon_pos + 1).c_str());
    return true;
}

void print_usage() {
    std::cerr << "Usage: udp_forward" << " -i <listen_ip:port> -o <target_ip:port>" << std::endl;
    std::cerr << "       udp_forward" <<  " --input=<listen_ip:port> --output=<target_ip:port>" << std::endl;    
}

int main(int argc, char *argv[]) {
    std::string listen_ip="0,0,0,0", target_ip="127.0.0.1";
    int listen_port = 5601, target_port = 5600;

    // 定义长选项
    struct option long_options[] = {
        {"input", required_argument, nullptr, 'i'},
        {"output", required_argument, nullptr, 'o'},
        {nullptr, 0, nullptr, 0}
    };

    int option_index = 0;
    int c;
    while ((c = getopt_long(argc, argv, "i:o:", long_options, &option_index)) != -1) {
        switch (c) {
            case 'i':
                if (!parse_ip_port(optarg, listen_ip, listen_port)) {
                    print_usage();
                    return EXIT_FAILURE;
                }
                break;
            case 'o':
                if (!parse_ip_port(optarg, target_ip, target_port)) {
                    print_usage();
                    return EXIT_FAILURE;
                }
                break;
            default:
                print_usage();
                return EXIT_FAILURE;
        }
    }

    // 检查是否提供了所有必要的参数
    if (listen_ip.empty() || listen_port == 0 || target_ip.empty() || target_port == 0) {
        print_usage();
        return EXIT_FAILURE;
    }

    // 初始化监听地址
    struct sockaddr_in listen_addr, target_addr;
    memset(&listen_addr, 0, sizeof(listen_addr));
    listen_addr.sin_family = AF_INET;
    listen_addr.sin_port = htons(listen_port);
    if (inet_pton(AF_INET, listen_ip.c_str(), &listen_addr.sin_addr) <= 0) {
        std::cerr << "Invalid listen IP address: " << listen_ip << std::endl;
        return EXIT_FAILURE;
    }

    // 创建监听UDP套接字
    int listen_fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (listen_fd < 0) {
        perror(" Failed to create listen socket");
        return EXIT_FAILURE;
    }

    // 绑定监听地址
    if (bind(listen_fd, (struct sockaddr *)&listen_addr, sizeof(listen_addr)) < 0) {
        perror(" Failed to bind listen socket");
        close(listen_fd);
        return EXIT_FAILURE;
    }
    std::cout << " Listening on " << listen_ip << ":" << listen_port << std::endl;

    // 初始化目标地址
    memset(&target_addr, 0, sizeof(target_addr));
    target_addr.sin_family = AF_INET;
    target_addr.sin_port = htons(target_port);
    if (inet_pton(AF_INET, target_ip.c_str(), &target_addr.sin_addr) <= 0) {
        std::cerr << "Invalid target IP address: " << target_ip << std::endl;
        close(listen_fd);
        print_usage();
        return EXIT_FAILURE;
    }
    std::cout << " Sending to " << target_ip << ":" << target_port << std::endl;
    char buffer[BUFFER_SIZE];
    socklen_t addr_len = sizeof(listen_addr);
    ssize_t bytes_received;
    // 主循环
    while (true) {
        // 接收数据
        bytes_received = recvfrom(listen_fd, buffer, BUFFER_SIZE, 0, 
                                          (struct sockaddr *)&listen_addr, &addr_len);
        if (bytes_received < 0) {
            perror("Failed to receive data");
            continue;
        }

        // 打印接收到的数据（可选）
        //std::cout << "Received " << bytes_received << " bytes from " 
        //          << inet_ntoa(listen_addr.sin_addr) << ":" 
        //          << ntohs(listen_addr.sin_port) << std::endl;

        // 转发数据
        if (sendto(listen_fd, buffer, bytes_received, 0, 
                   (struct sockaddr *)&target_addr, sizeof(target_addr)) < 0) {
            perror("Failed to send data to target");
        }
        /* else {
            std::cout << "Forwarded " << bytes_received << " bytes to " 
                      << target_ip << ":" << target_port << std::endl;
        }*/
    }

    // 关闭套接字
    close(listen_fd);
    return EXIT_SUCCESS;
}